from construct import *


identifier = Struct(
    "signature" / Const(b"\x7fELF"),
    "elfclass" / Enum(Byte, ELFCLASSNONE=0, ELFCLASS32=1, ELFCLASS64=2),
    "encoding" / Enum(Byte, ELFDATANONE=0, LSB=1, MSB=2),
    "version" / Enum(Byte, EV_NONE=0, EV_CURRENT=1),
    "osabi"
    / Enum(
        Byte,
        ELFOSABI_NONE=0,
        ELFOSABI_HPUX=1,
        ELFOSABI_NETBSD=2,
        ELFOSABI_GNU=3,
        ELFOSABI_SOLARIS=6,
        ELFOSABI_AIX=7,
        ELFOSABI_IRIX=8,
        ELFOSABI_FREEBSD=9,
        ELFOSABI_TRU64=10,
        ELFOSABI_MODESTO=11,
        ELFOSABI_OPENBSD=12,
        ELFOSABI_OPENVMS=13,
        ELFOSABI_NSK=14,
        ELFOSABI_AROS=15,
        ELFOSABI_FENIXOS=16,
        ELFOSABI_CLOUDABI=17,
        ELFOSABI_OPENVOS=18,
    ),
    "abiversion" / Byte,
    Padding(7),
)


def program_header(ELFInt32, ELFInt64, is64bit=True):
    """Program header structure."""
    Addr = IfThenElse(is64bit, ELFInt64, ELFInt32)

    return Struct(
        "p_type"
        / Enum(
            ELFInt32,
            PT_NULL=0x00000000,
            PT_LOAD=0x00000001,
            PT_DYNAMIC=0x00000002,
            PT_INTERP=0x00000003,
            PT_NOTE=0x00000004,
            PT_SHLIB=0x00000005,
            PT_PHDR=0x00000006,
            PT_TLS=0x00000007,
            PT_LDOOS=0x60000000,
            PT_HIOS=0x6FFFFFFF,
            PT_LOPROC=0x70000000,
            PT_HIPROC=0x7FFFFFFF,
        ),
        "flags_64"
        / If(
            is64bit,
            Enum(
                ELFInt32, R__=0x4, RW_=0x6, RWX=0x7, _WX=0x3, __X=0x1, _W_=0x2, R_X=0x5
            ),
        ),
        "offset" / Addr,
        "virtual_address" / Addr,
        "physical_address" / Addr,
        "size_file" / Addr,
        "size_mem" / Addr,
        "flags_32"
        / If(
            (not is64bit),
            Enum(
                ELFInt32, R__=0x4, RW_=0x6, RWX=0x7, _WX=0x3, __X=0x1, _W_=0x2, R_X=0x5
            ),
        ),
        "alignment" / Addr,
    )


def section_header(ELFInt32, ELFInt64, is64bit=True):
    """Section header structure."""
    Addr = IfThenElse(is64bit, ELFInt64, ELFInt32)

    return Struct(
        "sh_name_offset" / ELFInt32,
        "sh_name"
        / Pointer(this._.strtab_data_offset + this.sh_name_offset, CString("utf-8")),
        "sh_type"
        / Enum(
            ELFInt32,
            SHT_NULL=0x0,
            SHT_PROGBITS=0x1,
            SHT_SYMTAB=0x2,
            SHT_STRTAB=0x3,
            SHT_RELA=0x4,
            SHT_HASH=0x5,
            SHT_DYNAMIC=0x6,
            SHT_NOTE=0x7,
            SHT_NOBITS=0x7,
            SHT_REL=0x9,
            SHT_SHLIB=0x0A,
            SHT_DYNSYM=0x0B,
            SHT_INIT_ARRAY=0x0E,
            SHT_FINI_ARRAY=0x0F,
            SHT_PREINIT_ARRAY=0x10,
            SHT_GROUP=0x11,
            SHT_SYMTAB_SHNDX=0x12,
            SHT_NUM=0x13,
            SHT_LOOS=0x60000000,
        ),
        "sh_flags"
        / Enum(
            Addr,
            SHF_WRITE=0x1,
            SHF_ALLOC=0x2,
            SHF_EXECINSTR=0x4,
            SHF_MERGE=0x10,
            SHF_STRINGS=0x20,
            SHF_INFO_LINK=0x40,
            SHF_LINK_ORDER=0x80,
            SHF_OS_NONCONFORMING=0x100,
            SHF_GROUP=0x200,
            SHF_TLS=0x400,
            SHF_MASKOS=0x0FF00000,
            SHF_MASKPROC=0xF0000000,
            SHF_ORDERED=0x4000000,
            SHF_EXCLUDE=0x8000000,
        ),
        "sh_addr" / Addr,
        "sh_offset" / Addr,
        "sh_size" / Addr,
        "sh_link" / ELFInt32,
        "sh_info" / ELFInt32,
        "sh_addralign" / Addr,
        "sh_entsize" / Addr,
    )


def body(ELFInt16, ELFInt32, ELFInt64, is64bit=True):
    """ELF file structure."""
    p_header = program_header(ELFInt32, ELFInt64, is64bit)
    s_header = section_header(ELFInt32, ELFInt64, is64bit)

    return Struct(
        "type"
        / Enum(
            ELFInt16,
            ET_NONE=0,
            ET_REL=1,
            ET_EXEC=2,
            ET_DYN=3,
            ET_CORE=4,
            ET_LOOS=0xFE00,
            ET_HIOS=0xFEFF,
            ET_LOPROC=0xFF00,
            ET_HIPROC=0xFFFF,
        ),
        "machine"
        / Enum(
            ELFInt16,
            EM_NONE=0,
            EM_M32=1,
            EM_SPARC=2,
            EM_386=3,
            EM_68K=4,
            EM_88K=5,
            EM_860=7,
            EM_MIPS=8,
            EM_S370=9,
            EM_MIPS_RS3_LE=10,
            EM_PARISC=15,
            EM_VPP500=17,
            EM_SPARC32PLUS=18,
            EM_960=19,
            EM_PPC=20,
            EM_PPC64=21,
            EM_S390=22,
            EM_V800=36,
            EM_FR20=37,
            EM_RH32=38,
            EM_RCE=39,
            EM_ARM=40,
            EM_ALPHA=41,
            EM_SH=42,
            EM_SPARCV9=43,
            EM_TRICORE=44,
            EM_ARC=45,
            EM_H8_300=46,
            EM_H8_300H=47,
            EM_H8S=48,
            EM_H8_500=49,
            EM_IA_64=50,
            EM_MIPS_X=51,
            EM_COLDFIRE=52,
            EM_68HC12=53,
            EM_MMA=54,
            EM_PCP=55,
            EM_NCPU=56,
            EM_NDR1=57,
            EM_STARCORE=58,
            EM_ME16=59,
            EM_ST100=60,
            EM_TINYJ=61,
            EM_X86_64=62,
            EM_PDSP=63,
            EM_PDP10=64,
            EM_PDP11=65,
            EM_FX66=66,
            EM_ST9PLUS=67,
            EM_ST7=68,
            EM_68HC16=69,
            EM_68HC11=70,
            EM_68HC08=71,
            EM_68HC05=72,
            EM_SVX=73,
            EM_ST19=74,
            EM_VAX=75,
            EM_CRIS=76,
            EM_JAVELIN=77,
            EM_FIREPATH=78,
            EM_ZSP=79,
            EM_MMIX=80,
            EM_HUANY=81,
            EM_PRISM=82,
            EM_AVR=83,
            EM_FR30=84,
            EM_D10V=85,
            EM_D30V=86,
            EM_V850=87,
            EM_M32R=88,
            EM_MN10300=89,
            EM_MN10200=90,
            EM_PJ=91,
            EM_OPENRISC=92,
            EM_ARC_COMPACT=93,
            EM_XTENSA=94,
            EM_VIDEOCORE=95,
            EM_TMM_GPP=96,
            EM_NS32K=97,
            EM_TPC=98,
            EM_SNP1K=99,
            EM_ST200=100,
            EM_IP2K=101,
            EM_MAX=102,
            EM_CR=103,
            EM_F2MC16=104,
            EM_MSP430=105,
            EM_BLACKFIN=106,
            EM_SE_C33=107,
            EM_SEP=108,
            EM_ARCA=109,
            EM_UNICORE=110,
            EM_EXCESS=111,
            EM_DXP=112,
            EM_ALTERA_NIOS2=113,
            EM_CRX=114,
            EM_XGATE=115,
            EM_C166=116,
            EM_M16C=117,
            EM_DSPIC30F=118,
            EM_CE=119,
            EM_M32C=120,
            EM_TSK3000=131,
            EM_RS08=132,
            EM_SHARC=133,
            EM_ECOG2=134,
            EM_SCORE7=135,
            EM_DSP24=136,
            EM_VIDEOCORE3=137,
            EM_LATTICEMICO32=138,
            EM_SE_C17=139,
            EM_TI_C6000=140,
            EM_TI_C2000=141,
            EM_TI_C5500=142,
            EM_TI_ARP32=143,
            EM_TI_PRU=144,
            EM_MMDSP_PLUS=160,
            EM_CYPRESS_M8C=161,
            EM_R32C=162,
            EM_TRIMEDIA=163,
            EM_HEXAGON=164,
            EM_8051=165,
            EM_STXP7X=166,
            EM_NDS32=167,
            EM_ECOG1=168,
            EM_MAXQ30=169,
            EM_XIMO16=170,
            EM_MANIK=171,
            EM_CRAYNV2=172,
            EM_RX=173,
            EM_METAG=174,
            EM_MCST_ELBRUS=175,
            EM_ECOG16=176,
            EM_CR16=177,
            EM_ETPU=178,
            EM_SLE9X=179,
            EM_L10M=180,
            EM_K10M=181,
            EM_AARCH64=183,
            EM_AVR32=185,
            EM_STM8=186,
            EM_TILE64=187,
            EM_TILEPRO=188,
            EM_CUDA=190,
            EM_TILEGX=191,
            EM_CLOUDSHIELD=192,
            EM_COREA_1ST=193,
            EM_COREA_2ND=194,
            EM_ARC_COMPACT2=195,
            EM_OPEN8=196,
            EM_RL78=197,
            EM_VIDEOCORE5=198,
            EM_78KOR=199,
            EM_56800EX=200,
            EM_BA1=201,
            EM_BA2=202,
            EM_XCORE=203,
            EM_MCHP_PIC=204,
            EM_INTEL205=205,
            EM_INTEL206=206,
            EM_INTEL207=207,
            EM_INTEL208=208,
            EM_INTEL209=209,
            EM_KM32=210,
            EM_KMX32=211,
            EM_KMX16=212,
            EM_KMX8=213,
            EM_KVARC=214,
            EM_CDP=215,
            EM_COGE=216,
            EM_COOL=217,
            EM_NORC=218,
            EM_CSR_KALIMBA=219,
            EM_Z80=220,
            EM_VISIUM=221,
            EM_FT32=222,
            EM_MOXIE=223,
            EM_AMDGPU=224,
            EM_RISCV=243,
        ),
        "version" / Enum(ELFInt32, EV_NONE=0, EV_CURRENT=1),
        "entry" / IfThenElse(is64bit, ELFInt64, ELFInt32),
        "ph_offset" / IfThenElse(is64bit, ELFInt64, ELFInt32),
        "sh_offset" / IfThenElse(is64bit, ELFInt64, ELFInt32),
        "flags" / ELFInt32,
        "header_size" / ELFInt16,
        "ph_entry_size" / ELFInt16,
        "ph_count" / ELFInt16,
        "sh_entry_size" / ELFInt16,
        "sh_count" / ELFInt16,
        "strtab_section_index" / ELFInt16,
        "strtab_data_offset"
        / Pointer(
            this.sh_offset
            + this.strtab_section_index * this.sh_entry_size
            + (24 if is64bit else 16),
            ELFInt32,
        ),
        "program_table" / Pointer(this.ph_offset, p_header[this.ph_count]),
        "sections" / Pointer(this.sh_offset, s_header[this.sh_count]),
    )


# ELF file structure
elf = Struct(
    "identifier" / identifier,
    "body"
    / IfThenElse(
        this.identifier.encoding == "LSB",
        IfThenElse(
            this.identifier.elfclass == "ELFCLASS64",
            body(Int16ul, Int32ul, Int64ul, is64bit=True),
            body(Int16ul, Int32ul, Int64ul, is64bit=False),
        ),
        IfThenElse(
            this.identifier.elfclass == "ELFCLASS64",
            body(Int16ub, Int32ub, Int64ub, is64bit=True),
            body(Int16ub, Int32ub, Int64ub, is64bit=False),
        ),
    ),
)
